# !/usr/bin/env python
# http://code.google.com/p/beaconbits/
#
# This is the collector, it can be used with a coninuous collection from a flow feed or
# it can be use to read a file or file(s) saved in the argus flow format.
#
# code to parse flows from argus although any flow tool that can output comma seperated fields can be used.
# Seek specifically TCP packets, the associated time in unix date format,
# source IP along with Destination IP and port.

# source feed does not have to be Argus, any flow will work if you modify the properties correctly.
# The idea is to get the appropriate files into a DB for time series analysis.

# This code is released under license GPL2
# http://www.gnu.org/licenses/gpl-2.0.html


import sys
import subprocess
import argparse
import time
import redis

options = None

def main():
        global options
        parser = argparse.ArgumentParser(description='Argus flow read to redis.')
        ''' File(s) or interface '''

        # files require the designation plus the path
        parser.add_argument('-f','--file',help='specifiy the file by path /path/to/file')
        # Interface requies a server and port parsed to work effectively
        parser.add_argument('-i', '--interface', help='specify remote argus and optional port number  server:port')
        args = parser.parse_args()
        print args

        if args.file == None and args.interface == None:
            print 'you must select file or interface'
            print '-f /path/to/argus/file'
            print '-i 127.0.0.1:561 and requires that argus is serving flows on that port'
            sys.exit(-1)

        # process the parsed arguments into a valid argus command
        if args.file != None:
            command = "/usr/sbin/ra -nnr "+args.file+"  -c, -u -s stime saddr daddr dport - tcp and src syn 2> /dev/null"

        if args.interface != None:
            command = "/usr/sbin/ra -S ",args.interface,"  -c, -u -s stime saddr daddr dport - tcp and src syn 2> /dev/null"
        
        '''
        Parse fields to get the time in UNIX format followed by the source,destination and destination port
        1333181272.149791,1.1.1.1,192.168.1.1,80
        0	frame.time
        1	ip.src
        2       ip.dst
        3       dport
        '''

        ''' Calling ra (read argus) as a subprocess'''
        # Execute the argus/ra (read argus) process as given above
        argus = subprocess.Popen(command,stdout=subprocess.PIPE,shell=True)    

        #open a connection to the local redis database and assign to r
        r = redis.StrictRedis(host='localhost', port=6379, db=0)
        
        while True:
                argus.poll()                   # Constantly polling the argus/ra subprocess
                line = argus.stdout.readline() # Check for stdout
                if (line==''):
                    # if we're done processing a file or some sort of halt, gracefully exit
                    print "Finished."
                    sys.exit(1)
                else:
                    # strip the end of the line
                    line = line[:-1]
                    # set delimiter
                    fields = line.split(',')
                    if fields[0] == 'StartTime':
                        print 'header'
                    else:
                        # assign IPv4 source IP
                        ip_src = fields[1]
                        # assign IPv4 destination IP
                        ip_dst = fields[2]
                        # assign destination port
                        ip_dport = fields[3]
                        
                        # parse / quantize time to the second, not sub-second
                        Fixtime = int(line[0:10])
                        
                        # Shorten the UNIX time to seconds to our working_date
                        # prob should be a function
                        UDate = time.gmtime(Fixtime)
                        working_date = str(UDate.tm_year)+str(UDate.tm_mon)+str(UDate.tm_mday)

                try:
                    # quantset is concatenation of SET:1.1.1.11:192.168.1.1:80
                    # Store these sets in redis but first extract strings for ip source, ip dest, and dest port
                    # store sends in two keys, one for the specific date and one for multiple days (multi)
                    quantset_multi = str(ip_src)+":"+str(ip_dst)+":"+str(ip_dport)+":multi"
                    quantset = str(ip_src)+":"+str(ip_dst)+":"+str(ip_dport)+":"+working_date

                    # assign the keys for IP source and IP destination
                    keyset_ipsrc = "ip_src:" + str(ip_src)
                    keyset_ipdst = "ip_dst:" + str(ip_dst)

                    # Store the value of time to the set
                    r.sadd('SET:'+quantset, Fixtime)
                    r.sadd('SET:'+quantset_multi, Fixtime)

                    # tracking by count instances of source and destination IPs
                    r.incr(keyset_ipsrc)
                    r.incr(keyset_ipdst)

                # conditions for failure if needed and probably in the initial polling
                except:
                    print "Initial polling, error if this message repeats",line
  
if __name__ == '__main__':
        main()



